/*
 * Copyright [2013-2021], Alibaba Group Holding Limited
 *
 * Licensed under the Apache License, Version 2.0 (the "License");
 * you may not use this file except in compliance with the License.
 * You may obtain a copy of the License at
 *
 *    http://www.apache.org/licenses/LICENSE-2.0
 *
 * Unless required by applicable law or agreed to in writing, software
 * distributed under the License is distributed on an "AS IS" BASIS,
 * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 * See the License for the specific language governing permissions and
 * limitations under the License.
 */

package com.alibaba.polardbx.matrix.jdbc;

import com.alibaba.polardbx.common.logical.ITStatement;
import com.alibaba.polardbx.common.utils.GeneralUtil;
import com.alibaba.polardbx.common.utils.logger.Logger;
import com.alibaba.polardbx.common.utils.logger.LoggerFactory;
import com.alibaba.polardbx.druid.sql.parser.ByteString;
import com.alibaba.polardbx.optimizer.context.ExecutionContext;

import java.sql.ResultSet;
import java.sql.SQLException;

public class TStatement implements ITStatement {

    private static final Logger log = LoggerFactory.getLogger(TStatement.class);

    protected ByteString sql;
    protected TDataSource ds;
    protected TConnection conn;

    protected ExecutionContext executionContext = null;
    /**
     * 更新计数，如果执行了多次，那么这个值只会返回最后一次执行的结果。 如果是一个query，那么返回的数据应该是-1
     */
    protected int updateCount = -1;

    /**
     * 经过计算后的结果集，允许使用 getResult函数调用. 一个statement只允许有一个结果集
     */
    protected ResultSet currentResultSet;

    /**
     * 当前statment 是否是关闭的
     */
    protected boolean closed;

    public TStatement(TDataSource ds, TConnection tConnection, ExecutionContext executionContext) {
        this.ds = ds;
        this.conn = tConnection;
        this.executionContext = executionContext;
    }

    public TStatement(TDataSource ds, TConnection tConnection, ByteString sql, ExecutionContext executionContext) {
        this.ds = ds;
        this.conn = tConnection;
        this.sql = sql;
        this.executionContext = executionContext;
    }

    protected void checkClosed() throws SQLException {
        if (closed) {
            throw new SQLException("No operations allowed after statement closed.");
        }
    }

    protected ResultSet executeQuery(ByteString sql) throws SQLException {
        checkClosed();
        ensureResultSetIsEmpty();
        executeSQL(sql);
        return currentResultSet;
    }

    @Override
    public ResultSet executeQuery(String sql) throws SQLException {
        // Keep compatible with internal usage e.g. DDLNewHandler
        return executeQuery(ByteString.from(sql));
    }

    @Override
    public int executeUpdate(String sql) throws SQLException {
        return executeUpdateInternal(ByteString.from(sql), -1, null, null);
    }

    protected int executeUpdateInternal(ByteString sql, int autoGeneratedKeys, int[] columnIndexes,
                                        String[] columnNames) throws SQLException {
        checkClosed();
        ensureResultSetIsEmpty();
        executionContext.setAutoGeneratedKeys(autoGeneratedKeys);
        executionContext.setColumnIndexes(columnIndexes);
        executionContext.setColumnNames(columnNames);
        executeSQL(sql);
        if (currentResultSet instanceof TMultiResultSet) {
            return ((TMultiResultSet) currentResultSet).getAffectRows();
        } else {
            return ((TResultSet) currentResultSet).getAffectRows();
        }
    }

    protected void executeSQL(ByteString sql) throws SQLException {
        currentResultSet = conn.executeSQL(sql, null, this, executionContext);
    }

    @Override
    public void close() throws SQLException {
        close(true);
    }

    public void close(boolean removeThis) throws SQLException {
        close(removeThis, true);
    }

    public void close(boolean removeStmt, boolean closeResultSet) throws SQLException {
        if (closed) {
            return;
        }
        try {
            if (currentResultSet != null && closeResultSet) {
                currentResultSet.close();
            }

            if (removeStmt) {
                conn.removeStatement(this);
            }

            if (executionContext.isInternalSystemSql()) {
                // For all stmt generated by drds internal system(e.g. sync
                // xxx/async ddl jobs/ and so on),
                // all the memory pool of these sql should be destroy here
                executionContext.clearAllMemoryPool();
            }
        } catch (Exception e) {
            throw GeneralUtil.nestedException(e);
        } finally {
            currentResultSet = null;
        }
        closed = true;
    }

    /**
     * 如果新建了查询，那么上一次查询的结果集应该被显示的关闭掉。这才是符合jdbc规范的
     */
    protected void ensureResultSetIsEmpty() throws SQLException {
        if (currentResultSet != null) {
            // log.debug("result set is not null,close current result set");
            try {
                currentResultSet.close();
            } catch (SQLException e) {
                log.error("exception on close last result set . can do nothing..", e);
            } finally {
                // 最终要显示的关闭它
                currentResultSet = null;
            }
        }
    }

    public boolean getMoreResults() throws SQLException {
        return getMoreResults(1);
    }

    private boolean getMoreResults(int current) throws SQLException {
        if (currentResultSet instanceof TMultiResultSet) {
            boolean more = ((TMultiResultSet) currentResultSet).getMoreResults(current);
            updateCount = ((TMultiResultSet) currentResultSet).getAffectRows();
            return more;
        } else {
            updateCount = -1;
            return false;
        }
    }
}
